import math


class GaussianMixture:
    """只支持1个特征的高斯混合模型能参数估计的EM算法"""

    def __init__(self, X, init_mean, n_components, max_iter=100):
        self.X = X
        self.n_components = n_components
        self.max_iter = max_iter

        self.n_samples = len(X)

        # 取参数的初始值
        self.alpha = [1 / self.n_components] * self.n_components  # 分模型权重
        self.means = init_mean  # 均值
        self.sigma = [1.0] * self.n_components  # 方差

        # 执行训练
        self._train()

    def _train(self):
        for _ in range(self.max_iter):
            # E步：根据当前模型参数，计算分模型对观测数据的响应度
            gamma = [[0] * self.n_components for _ in range(self.n_samples)]
            for j in range(self.n_samples):
                sum_ = 0
                for k in range(self.n_components):
                    gamma[j][k] = self.alpha[k] * self._count_gaussian(self.X[j], k)
                    sum_ += gamma[j][k]
                for k in range(self.n_components):
                    gamma[j][k] /= sum_

            # M步：计算新一轮迭代的模型参数
            means_new = [0.0] * self.n_components
            for k in range(self.n_components):
                sum1 = 0  # 分子
                sum2 = 0  # 分母
                for j in range(self.n_samples):
                    sum1 += gamma[j][k] * self.X[j]
                    sum2 += gamma[j][k]
                means_new[k] = sum1 / sum2

            sigma_new = [1.0] * self.n_components
            for k in range(self.n_components):
                sum1 = 0  # 分子
                sum2 = 0  # 分母
                for j in range(self.n_samples):
                    sum1 += gamma[j][k] * math.pow(self.X[j] - self.means[k], 2)
                    sum2 += gamma[j][k]
                sigma_new[k] = sum1 / sum2

            alpha_new = [1 / self.n_components] * self.n_components
            for k in range(self.n_components):
                sum1 = 0  # 分子
                for j in range(self.n_samples):
                    sum1 += gamma[j][k]
                alpha_new[k] = sum1 / self.n_samples

            self.alpha = alpha_new
            self.means = means_new
            self.sigma = sigma_new

    def _count_gaussian(self, x, k):
        """计算高斯密度函数"""
        return math.pow(math.e, -math.pow(x - self.means[k], 2) / (2 * self.sigma[k])) / (
            math.sqrt(2 * math.pi * self.sigma[k]))

    def predict(self, x):
        best_k, best_g = -1, 0
        for k in range(self.n_components):
            g = self.alpha[k] * self._count_gaussian(x, k)
            if g > best_g:
                best_k, best_g = k, g
        return best_k
